"""
target detection
"""
import rospy
import cv2
from std_msgs.msg import String
from sensor_msgs.msg import Image
from swarm_msgs.msg import BoundingBox, BoundingBoxes
import os
import re
import numpy as np
import platform
import time
from datetime import datetime
import threading

rospy.init_node('my_node_name', anonymous=True)
from YoloRknn import YoloRknn

yolo = YoloRknn('/home/orangepi/kbt_ws/src/kbtdj_air/det/scripts/1119.rknn')

class FrameReader(threading.Thread):
  def __init__(self):
    threading.Thread.__init__(self)
    self.cap = cv2.VideoCapture(0)  # 从摄像头读取视频
    self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 960)
    self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)
    self.frame = None
    self.running = True
    self.lock = threading.Lock()
    self.bbox_pub = rospy.Publisher("/intercept/horizonal_cam", BoundingBoxes, queue_size=10)
    if not self.cap.isOpened():
      raise ValueError("Could not open video device")

  def run(self):
    while self.running:
      self.cap.grab()
      ret, frame = self.cap.retrieve()
      with self.lock:
        if ret:
          self.frame = frame.copy()
        else:
          self.frame = None
      time.sleep(0.01)

  def stop(self):
    self.running = False
    self.cap.release()


class YoloProcessor(threading.Thread):
  def __init__(self, frame_reader):
    threading.Thread.__init__(self)
    self.frame_reader = frame_reader
    self.running = True
    self.write_count = 0
    self.bboxes = BoundingBoxes()
    self.org_img = np.zeros((552, 640, 3), dtype=np.uint8)  # 创建一个空图像

  def run(self):
    while self.running:
      frame = None
      with self.frame_reader.lock:
        if self.frame_reader.frame is not None:
          frame = self.frame_reader.frame.copy()
      if frame is not None:
        self.process_frame(frame)
      time.sleep(0.01)  # 控制处理帧的速度

  def process_frame(self, frame):
    try:
      self.bboxes.bounding_boxes = []
      yolo.det(frame)
      for det_bbox in yolo.det_bboxes:
        bbox = BoundingBox()
        bbox.xmin = det_bbox['left']
        bbox.ymin = det_bbox['top']
        bbox.xmax = det_bbox['right']
        bbox.ymax = det_bbox['bottom']
        bbox.Class = det_bbox['id']
        print(det_bbox['id'])
        if bbox.Class == 'uav':
          self.bboxes.bounding_boxes.append(bbox)
          print(bbox)
      if len(self.bboxes.bounding_boxes) > 0:
        try:
          self.frame_reader.bbox_pub.publish(self.bboxes)
          img_1 = cv2.resize(yolo.res_img, (640, 552))
          if img_1.size > 0:
            cv2.imshow("horizonal_cam window", img_1)
            cv2.moveWindow("horizonal_cam window", 1, 1)
            cv2.waitKey(1)
        except Exception as e:
          print(f"An 1error occurred: {e}")
      else:
        img_1 = cv2.resize(yolo.org_img, (640, 552))
        if img_1.size > 0:
          cv2.imshow("horizonal_cam window", img_1)
          cv2.moveWindow("horizonal_cam window", 1, 1)
          cv2.waitKey(1)


      #write image for train
      img_111 = cv2.resize(yolo.org_img, (640, 552))
      if img_111.size > 0:
          if self.write_count%5==0:
              save_dir = '/home/orangepi/pictures_yolo_horizonal'
              if not os.path.exists(save_dir):
                  os.makedirs(save_dir)
              now = datetime.now()
              timestamp = now.strftime('%Y%m%d_%H%M%S')
              filename = os.path.join(save_dir, f'{timestamp}.png')
              cv2.imwrite(filename, img_111)
          self.write_count+=1  


    except Exception as e:
      print(f"An 2error occurred: {e}")
      cv2.imshow("horizonal_cam window", frame)
      cv2.waitKey(1)

  def stop(self):
    self.running = False

def main():
  frame_reader = FrameReader()
  yolo_processor = YoloProcessor(frame_reader)

  frame_reader.start()
  yolo_processor.start()
  try:
    print("1")
    # while True:
    #     frame = None
    #     with frame_reader.lock:
    #         if frame_reader.frame is not None:
    #             frame = frame_reader.frame.copy()
    #
    #     if frame is not None:
    #         cv2.imshow('Frame', frame)
    #         if cv2.waitKey(1) & 0xFF == ord('q'):
    #             break
    #     time.sleep(0.01)
  except KeyboardInterrupt:
    frame_reader.stop()
    yolo_processor.stop()
  except Exception as e:
    print(f"An error occurred in main: {e}")

  frame_reader.join()
  yolo_processor.join()
  cv2.destroyAllWindows()


if __name__ == '__main__':
  main()
  rospy.spin()
